# -*- coding: utf-8 -*-
"""
@Time    : 2024/7/29 14:17 
@Author  : ZhangShenao 
@File    : 4.实现可动态配置的Retriever.py
@Desc    : 实现可动态配置的Retriever

BaseRetriever是所有Retriever组件的基类
它继承了RunnableSerializable,因此它是一个Runnable可运行组件,可以直接通过LCEL表达式执行
也可以通过.configurable_fields()方法,动态配置运行时参数
"""

import dotenv
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_core.runnables import ConfigurableField
from langchain_pinecone import PineconeVectorStore

# 加载环境变量
dotenv.load_dotenv()

embeddings = ZhipuAIEmbeddings(model='embedding-2')
vector_store = PineconeVectorStore(
    index_name='dataset',  # 指定索引名称
    embedding=embeddings,  # 指定Embedding模型
    namespace='llm-ops',  # 指定Namespace
    text_key='text',  # 指定文本属性key
)

# Retrieving文档检索
# retriever是Runnable可运行组件,可以通过LCEL表达式直接运行
retriever = vector_store.as_retriever(
    search_type='similarity_score_threshold',
    # 指定检索参数
    search_kwargs={
        'k': 3,
        'score_threshold': 0.5,
    }
).configurable_fields(
    # 指定可动态配置的参数
    search_type=ConfigurableField(id='search_type'),
    search_kwargs=ConfigurableField(id='search_kwargs'),
)

# 执行Chain,并为可配置字段动态赋值
result = retriever.with_config(configurable={
    'search_type': 'mmr',
    'search_kwargs': {
        'k': 2,
        'fetch_k': 20,
        'lambda_mult': 0.5,
    }
}).invoke('骆驼祥子和刘四爷是什么关系？')
print(result)
